"""
基于 LangChain 的数据收集智能体
"""
from langchain_core.tools import BaseTool
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from service.news_kimi_das import get_sina_finance_news
from service.tushare_pool import getDayStock
from agents.base_agent import StockAgent
import json
from datetime import datetime, timedelta
import pandas as pd
from colorama import Fore, Style
from typing import Dict, Any, Optional
from pydantic import BaseModel, Field


class StockDataTool(BaseTool):
    """股票数据获取工具"""
    name: str = "stock_data_tool"
    description: str = "获取股票的历史价格数据和基本信息"
    
    def _run(self, stock_code: str, days: int = 365) -> str:
        """获取股票数据"""
        try:
            today = datetime.now()
            start_date = (today - timedelta(days=days)).strftime("%Y%m%d")
            end_date = today.strftime("%Y%m%d")
            
            daily_k_data = getDayStock(
                stockCode=stock_code, 
                startDate=start_date, 
                endDate=end_date
            )
            
            if daily_k_data is None or daily_k_data.empty:
                return f"无法获取股票代码 {stock_code} 的数据"
            
            # 返回最近10条数据的摘要
            recent_data = daily_k_data.head(10)
            return f"股票 {stock_code} 最近10个交易日数据:\n{recent_data.to_string()}"
            
        except Exception as e:
            return f"获取股票数据时出错: {str(e)}"


class NewsSearchTool(BaseTool):
    """新闻搜索工具"""
    name: str = "news_search_tool"
    description: str = "搜索与股票相关的最新财经新闻"
    
    def _run(self, stock_code: str, query: Optional[str] = None) -> str:
        """搜索新闻"""
        try:
            if query is None:
                query = f"搜索关于股票 {stock_code} 的最新财经新闻"
            
            # 打印开始搜索日志
            print(Fore.CYAN + f"    - [NewsSearchTool] 开始搜索股票 {stock_code} 的新闻...")
            print(Fore.CYAN + f"    - [NewsSearchTool] 查询内容: {query}")
            
            # 调用 get_sina_finance_news 获取新闻
            news_result = get_sina_finance_news(query=query)
            
            # 打印获取结果日志
            print(Fore.GREEN + f"    - [NewsSearchTool] 成功获取新闻数据")
            
            # 尝试解析JSON格式的新闻数据
            try:
                news_data = json.loads(news_result)
                if isinstance(news_data, dict) and "news_list" in news_data:
                    news_list = news_data["news_list"]
                    print(Fore.GREEN + f"    - [NewsSearchTool] 解析到 {len(news_list)} 条新闻")
                    
                    # 格式化新闻输出
                    formatted_news = f"关于 {stock_code} 的最新新闻:\n"
                    formatted_news += "=" * 50 + "\n"
                    for idx, news in enumerate(news_list, 1):
                        formatted_news += f"\n{idx}. {news.get('title', '无标题')}\n"
                        formatted_news += f"   时间: {news.get('time', '未知')}\n"
                        formatted_news += f"   来源: {news.get('source', '未知')}\n"
                        formatted_news += f"   内容: {news.get('content', '无内容')}\n"
                        if news.get('url'):
                            formatted_news += f"   链接: {news.get('url')}\n"
                        formatted_news += "-" * 40 + "\n"
                    
                    return formatted_news
                else:
                    # 如果不是预期的JSON格式，直接返回原始内容
                    print(Fore.YELLOW + "    - [NewsSearchTool] 新闻数据非标准JSON格式，返回原始内容")
                    return f"关于 {stock_code} 的最新新闻:\n{news_result}"
            except json.JSONDecodeError:
                # 如果无法解析为JSON，直接返回原始内容
                print(Fore.YELLOW + "    - [NewsSearchTool] 无法解析为JSON，返回原始内容")
                return f"关于 {stock_code} 的最新新闻:\n{news_result}"
            
        except Exception as e:
            error_msg = f"搜索新闻时出错: {str(e)}"
            print(Fore.RED + f"    - [NewsSearchTool] {error_msg}")
            return error_msg


class DataCollectorAgent(StockAgent):
    """基于 LangChain 的数据收集智能体"""
    
    def __init__(self):
        super().__init__("data_collector")
        self.tools = self._create_tools()
        self.data_collection_chain = self._create_data_collection_chain()

    def _create_tools(self):
        """创建工具集"""
        return {
            "stock_data": StockDataTool(),
            "news_search": NewsSearchTool()
        }

    def _create_system_prompt(self) -> str:
        """重写系统提示词"""
        return """你是一个专业的股票数据收集师。你的任务是：
        1. 收集指定股票的历史价格数据
        2. 搜索相关的财经新闻
        3. 整理数据并生成结构化的报告

        请确保数据的准确性和完整性，并以清晰的格式呈现收集到的信息。"""

    def _create_data_collection_chain(self):
        """创建数据收集处理链"""
        
        def collect_stock_data(inputs: Dict[str, Any]) -> Dict[str, Any]:
            """收集股票数据"""
            stock_code = inputs["stock_code"]
            
            print(Fore.YELLOW + "    - [LangChain Tool] 正在获取股票数据...")
            stock_data = self.tools["stock_data"]._run(stock_code)
            
            # 打印股票数据结果
            print(Fore.BLUE + "\n" + "=" * 60)
            print(Fore.BLUE + "    📊 股票数据获取结果:")
            print(Fore.WHITE + "-" * 60)
            print(stock_data[:500] if len(stock_data) > 500 else stock_data)  # 限制打印长度
            if len(stock_data) > 500:
                print(Fore.YELLOW + f"    ... (数据过长，已截断，总长度: {len(stock_data)} 字符)")
            print(Fore.BLUE + "=" * 60 + "\n")
            
            print(Fore.YELLOW + "    - [LangChain Tool] 正在搜索相关新闻...")
            news_data = self.tools["news_search"]._run(stock_code)
            
            # 打印新闻数据结果
            print(Fore.CYAN + "\n" + "=" * 60)
            print(Fore.CYAN + "    📰 新闻数据获取结果:")
            print(Fore.WHITE + "-" * 60)
            print(news_data[:800] if len(news_data) > 800 else news_data)  # 新闻可能更长，允许更多内容
            if len(news_data) > 800:
                print(Fore.YELLOW + f"    ... (数据过长，已截断，总长度: {len(news_data)} 字符)")
            print(Fore.CYAN + "=" * 60 + "\n")
            
            return {
                "stock_code": stock_code,
                "stock_data": stock_data,
                "news_data": news_data
            }
        
        def format_report(data: Dict[str, Any]) -> str:
            """格式化报告"""
            report_template = """
            # 股票 {stock_code} 数据收集报告

            ## 1. 股票数据
            {stock_data}

            ## 2. 相关新闻
            {news_data}

            ## 3. 数据收集总结
            数据收集时间: {timestamp}
            数据来源: Tushare API, 财经新闻API
            数据状态: 收集完成
            """
            
            return report_template.format(
                stock_code=data["stock_code"],
                stock_data=data["stock_data"],
                news_data=data["news_data"],
                timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            )
        
        # 创建处理链
        chain = (
            RunnableLambda(collect_stock_data)
            | RunnableLambda(format_report)
        )
        
        return chain

    def execute(self, stock_code: str) -> str:
        """执行数据收集任务"""
        try:
            print(Fore.MAGENTA + f"\n    - [LangChain] 启动数据收集链处理 {stock_code}...")
            
            # 使用 LangChain 链处理
            result = self.data_collection_chain.invoke({"stock_code": stock_code})
            
            print(Fore.GREEN + "    - [LangChain] 数据收集链处理完成")
            return result
            
        except Exception as e:
            error_msg = f"数据收集过程中发生错误: {str(e)}"
            print(Fore.RED + f"    - [Error] {error_msg}")
            return error_msg
